import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont

# 读取图像
image_path = r'D:\Users\Administrator\Desktop\ipc_control\file\test_image\SEG1_CJ.jpg'
image = cv2.imread(image_path)

# 显示图像
# cv2.imshow('image', image)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

# 将图像转换为灰度图
gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

# (5, 5) 是高斯核的大小，可以根据需要调整
blurred_image = cv2.GaussianBlur(gray_image, (7, 7), 0)

# 计算图像的平均值
mean_value = np.mean(blurred_image) * 2

# 使用平均值作为阈值进行二值化
_, binary_image = cv2.threshold(blurred_image, mean_value, 255, cv2.THRESH_BINARY)

# 计算白色和黑色区域的像素数量
white_area = np.sum(binary_image == 255)
black_area = np.sum(binary_image == 0)

# 如果白色区域多，则反转图像
if white_area > black_area:
    binary_image = cv2.bitwise_not(binary_image)

# 创建一个结构元素（矩形形态）
kernel = np.ones((19, 19), np.uint8)  # 5x5的矩形结构元素

# 膨胀操作
dilated_image = cv2.dilate(binary_image, kernel, iterations=1)  # iterations=1表示执行一次膨胀操作

# 腐蚀操作
eroded_image = cv2.erode(dilated_image, kernel, iterations=1)  # iterations=1表示执行一次腐蚀操作

dilated_colored_image = cv2.cvtColor(eroded_image, cv2.COLOR_GRAY2BGR)

# 显示图像
# cv2.imshow('dilated_colored_image', dilated_colored_image)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

import paddlehub as hub

# 初始化 OCR 模块
ocr = hub.Module(name="chinese_ocr_db_crnn_mobile", enable_mkldnn=True)  # mkldnn加速仅在CPU下有效

# 识别文字
result = ocr.recognize_text(images=[dilated_colored_image])

# 转换为 Pillow 图像
image_pil = Image.fromarray(cv2.cvtColor(dilated_colored_image, cv2.COLOR_BGR2RGB))

# 创建绘图对象
draw = ImageDraw.Draw(image_pil)

# 加载中文字体文件（请确保路径正确）
font_path = "C:/Windows/Fonts/simsun.ttc"  # Windows系统的常用字体
font = ImageFont.truetype(font_path, 20)  # 设置字体大小

# 在图像上标记文字位置和内容
for item in result[0]['data']:
    text = item['text']  # 识别的文字内容
    points = item['text_box_position']  # 文本框的四个顶点坐标
    points = np.array(points, dtype=np.int32)  # 转换为numpy数组

    # 调试：打印每个文本框的坐标
    print("Points:", points)
    print("text:", text)

    # 使用Pillow绘制矩形框
    draw.line([tuple(points[0]), tuple(points[1]), tuple(points[2]), tuple(points[3]), tuple(points[0])],
              fill=(0, 255, 0), width=2)  # 绘制绿色矩形框

    # 计算左上角点的位置，用于显示文字
    top_left = (points[0][0], points[0][1] - 25)  # 上移一点，避免遮挡框线

    # 使用Pillow在图像上标记识别到的文字
    draw.text(top_left, text, font=font, fill=(255, 0, 0))  # 红色文字

# 将 Pillow 图像转换回 OpenCV 格式
dilated_colored_image = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)

# 显示图像
cv2.imshow('image', image)
cv2.imshow('Marked Image', dilated_colored_image)
cv2.waitKey(0)
cv2.destroyAllWindows()
